package com.framework.websocket.core;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.framework.websocket.model.SubscribeMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.Session;
import java.util.Map;
import java.util.HashMap;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket 连接管理器
 * 统一管理所有 WebSocket 连接和消息发送
 */
@Slf4j
@Component
public class WebSocketManager {
    
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    /**
     * 存储所有活跃的 WebSocket 会话
     */
    private final Map<String, Session> sessions = new ConcurrentHashMap<>();
    
    /**
     * 会话ID到客户端ID的映射
     */
    private final Map<String, String> sessionToClientId = new ConcurrentHashMap<>();
    
    /**
     * 添加新连接
     */
    public String addSession(Session session) {
        String clientId = UUID.randomUUID().toString();
        String sessionId = session.getId();
        
        sessions.put(sessionId, session);
        sessionToClientId.put(sessionId, clientId);
        
        log.info("新的 WebSocket 连接建立 - SessionId: {}, ClientId: {}", sessionId, clientId);
        return clientId;
    }
    
    /**
     * 移除连接
     */
    public void removeSession(Session session) {
        String sessionId = session.getId();
        String clientId = sessionToClientId.remove(sessionId);
        sessions.remove(sessionId);
        
        log.info("WebSocket 连接断开 - SessionId: {}, ClientId: {}", sessionId, clientId);
    }
    
    /**
     * 获取会话
     */
    public Session getSession(String sessionId) {
        return sessions.get(sessionId);
    }
    
    /**
     * 获取客户端ID
     */
    public String getClientId(String sessionId) {
        return sessionToClientId.get(sessionId);
    }
    
    /**
     * 发送消息给指定会话
     */
    public void sendMessage(Session session, String command, String params, Object response) {
        if (session == null || !session.isOpen()) {
            return;
        }
        
        try {
            SubscribeMessage message = new SubscribeMessage();
            message.setCommand(command);
            message.setParams(params);
            message.setResponse(response);
            message.setTimestamp(System.currentTimeMillis());
            
            String jsonMessage = objectMapper.writeValueAsString(message);
            session.getAsyncRemote().sendText(jsonMessage);
            
            log.debug("发送消息给会话 {} - Command: {}", session.getId(), command);
        } catch (Exception e) {
            log.error("发送消息失败 - SessionId: {}", session.getId(), e);
        }
    }
    
    /**
     * 发送推送消息给指定会话（新方法）
     */
    public void sendPushMessage(Session session, String topic, String key, Object data) {
        if (session == null || !session.isOpen()) {
            return;
        }
        
        try {
            SubscribeMessage message = new SubscribeMessage();
            message.setCommand("PUSH");
            message.setTopic(topic);
            message.setKey(key);
            message.setResponse(data);
            message.setTimestamp(System.currentTimeMillis());
            
            String jsonMessage = objectMapper.writeValueAsString(message);
            session.getAsyncRemote().sendText(jsonMessage);
            
            log.debug("推送消息给会话 {} - Topic: {}, Key: {}", session.getId(), topic, key);
        } catch (Exception e) {
            log.error("推送消息失败 - SessionId: {}", session.getId(), e);
        }
    }
    
    /**
     * 广播消息给所有连接
     */
    public void broadcast(String command, Object data) {
        sessions.values().parallelStream().forEach(session -> {
            sendMessage(session, command, null, data);
        });
        
        log.info("广播消息 - Command: {}, 连接数: {}", command, sessions.size());
    }
    
    /**
     * 发送错误消息
     */
    public void sendError(Session session, String error) {
        Map<String, Object> errorMap = new HashMap<>();
        errorMap.put("error", error);
        sendMessage(session, "ERROR", null, errorMap);
    }
    
    /**
     * 获取当前连接数
     */
    public int getConnectionCount() {
        return sessions.size();
    }
    
    /**
     * 获取所有连接信息
     */
    public Map<String, String> getConnectionInfo() {
        Map<String, String> info = new HashMap<>();
        sessionToClientId.forEach((sessionId, clientId) -> {
            Session session = sessions.get(sessionId);
            if (session != null) {
                info.put(clientId, session.isOpen() ? "connected" : "disconnected");
            }
        });
        return info;
    }
    
    /**
     * 检查会话是否存活
     */
    public boolean isSessionAlive(String sessionId) {
        Session session = sessions.get(sessionId);
        return session != null && session.isOpen();
    }
}